# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: xguo

import json
import logging
import os

from tornado.ioloop import IOLoop
from tornado import gen
from tornado.websocket import websocket_connect

from coin.base.timestamp import get_timestamp

BINANCE_WSS_URL = 'wss://stream.binance.com:9443/ws'
BINANCE_FUTURES_WSS_URL = 'wss://fstream.binance.com/ws'

if os.environ.get('USE_BINANCE_TESTNET'):
  BINANCE_FUTURES_WSS_URL = 'wss://testnet.binancefuture.com/ws'


class BinancePrivateSubscriber(object):
  def __init__(
      self,
      private_client,
      ioloop=None,
      notification_callback=None,
      close_callback=None,
      heartbeat_period_sec=300,
      ws_url=BINANCE_WSS_URL,
  ):
    self._private_client = private_client
    self._ioloop = ioloop or IOLoop.current()
    self._ws = None
    self._ws_url = ws_url
    self._ready = False

    self._notification_callback = notification_callback
    self._close_callback = close_callback

    self._listen_key = None
    self._last_received_message_timestamp = None
    self._heartbeat_period_sec = heartbeat_period_sec

    self._stopped = False
    self._init_params = (private_client,
                         ioloop,
                         notification_callback,
                         close_callback,
                         heartbeat_period_sec)
    self._logger = logging.getLogger(__name__)
    self._last_keep_alive_timestamp = get_timestamp()

  @property
  def last_received_message_timestamp(self):
    return self._last_received_message_timestamp

  def is_ready(self):
    return self._ready

  @gen.coroutine
  def _periodic_keep_alive(self):
    while True:
      yield gen.sleep(self._heartbeat_period_sec)
      if self._stopped:
        break

      if not self.is_ready():
        continue
      try:
        response = yield self._private_client.keep_alive_data_stream(self._listen_key)
        self._logger.info('keep alive: %s', response.status_code)
        if response.status_code == 200:
          self._last_keep_alive_timestamp = get_timestamp()
        else:
          self._check_last_keep_alive_timestamp()
      except Exception:
        self._logger.exception('Exception raised when put keep alive!')
        self._check_last_keep_alive_timestamp()

  def _check_last_keep_alive_timestamp(self):
    if get_timestamp() - self._last_keep_alive_timestamp > 660e9:
      self._logger.error('Keep alive failed twice in a row, close websocket!')
      self.stop()

  def start(self):
    self._ioloop.add_callback(self._connect)
    self._ioloop.add_callback(self._periodic_keep_alive)

  @gen.coroutine
  def _connect(self):
    # noinspection PyBroadException
    try:
      self._logger.info('%s try to connect to the websocket.', type(self).__name__)
      response = yield self._private_client.start_data_stream()
      self._listen_key = response.json()['listenKey']
      self._logger.info('listen_key: %s', self._listen_key)
      self._ws = yield websocket_connect('%s/%s' % (self._ws_url, self._listen_key))
      self._logger.info('Binance data stream websocket connected!')
      self._ioloop.add_callback(self._run)
      self._ready = True
    except Exception:
      self._logger.exception('%s connect to websocket failed!', type(self).__name__)
      self.stop()

  @gen.coroutine
  def _run(self):
    while not self._stopped:
      msg = None
      # noinspection PyBroadException
      try:
        msg = yield self._ws.read_message()
      except Exception:
        self._logger.exception('Error reading ws message')
        self.stop()

      if msg is None:
        self._logger.error('Websocket is closed!')
        self.stop()
        return

      if len(msg) == 0:
        self._logger.info('Empty message received!')
        continue

      # noinspection PyBroadException
      try:
        self.on_message(msg)
      except Exception:
        self._logger.exception('Exception raised when process message!')

  def stop(self):
    if self._stopped:
      return
    self._ioloop.add_callback(self._async_stop)

  @gen.coroutine
  def _async_stop(self):
    response = yield self._private_client.close_data_stream(self._listen_key)
    self._logger.info('Binance private subscriber (%s) stopped! %s', id(self), response.json())
    self._stopped = True
    self._ready = False
    self._ws.close()
    self._ws = None
    self._close_callback()

  def clone(self):
    return type(self)(*self._init_params)

  def on_message(self, raw_msg: bytes):
    self._last_received_message_timestamp = get_timestamp()
    if self._stopped:
      return

    try:
      data = json.loads(raw_msg)
    except ValueError:
      self._logger.exception('Invalid json string!, %s', raw_msg)
      return
    self._notification_callback(data)
